import re
import uuid
import queue
import hashlib
import base64
import secrets
import requests
import threading
from faker import Faker
from DrissionPage import Chromium
from helper.email import EmailServer

enable_register_log = True

class CursorRegister:
    CURSOR_URL = "https://www.cursor.com/"
    CURSOR_SIGNIN_URL = "https://authenticator.cursor.sh"
    CURSOR_PASSWORD_URL = "https://authenticator.cursor.sh/password"
    CURSOR_MAGAIC_CODE_URL = "https://authenticator.cursor.sh/magic-code"
    CURSOR_SIGNUP_URL = "https://authenticator.cursor.sh/sign-up"
    CURSOR_SIGNUP_PASSWORD_URL = "https://authenticator.cursor.sh/sign-up/password"
    CURSOR_EMAIL_VERIFICATION_URL = "https://authenticator.cursor.sh/email-verification"
    CURSOR_SETTING_URL = "https://www.cursor.com/settings"
    CURSOR_USAGE_URL = "https://www.cursor.com/api/usage"

    def __init__(self, 
                 browser: Chromium,
                 email_server: EmailServer = None):

        self.browser = browser
        self.email_server = email_server
        self.email_queue = queue.Queue()
        self.email_thread = None

        self.thread_id = threading.current_thread().ident
        self.retry_times = 5

    def sign_in(self, email, password = None):
        """注册Cursor账号"""
        assert any(x is not None for x in (self.email_server, password)), "应提供电子邮件服务器或密码，至少需要其中一项。"
 
        if self.email_server is not None:
            self.email_thread = threading.Thread(target=self.email_server.wait_for_new_message_thread,
                                                 args=(self.email_queue, ), 
                                                 daemon=True)
            self.email_thread.start()

        tab = self.browser.new_tab(self.CURSOR_SIGNIN_URL)
        # 输入邮箱
        for retry in range(self.retry_times):
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 输入邮箱")
                tab.ele("xpath=//input[@name='email']").input(email, clear=True)
                tab.ele("@type=submit").click()

                # 如果没有进入密码页面，尝试通过验证码
                if not tab.wait.url_change(self.CURSOR_PASSWORD_URL, timeout=3) and self.CURSOR_SIGNIN_URL in tab.url:
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 尝试通过邮箱页面的验证码")
                    self._cursor_turnstile(tab)

            except Exception as e:
                print(f"[注册][{self.thread_id}] 处理邮箱页面时出现异常")
                print(e)

            # 在密码页面或数据已验证，继续下一页
            if tab.wait.url_change(self.CURSOR_PASSWORD_URL, timeout=5):
                print(f"[注册][{self.thread_id}] 继续进入密码页面")
                break

            tab.refresh()
            # 超时退出函数
            if retry == self.retry_times - 1:
                print(f"[注册][{self.thread_id}] 输入邮箱地址超时")
                return tab, False

        # 检查是否需要注册新账号
        is_signup = False
        if tab.wait.eles_loaded("xpath=//button[contains(text(), '注册')]", timeout=1) or tab.wait.eles_loaded("xpath=//button[contains(text(), 'Sign up')]", timeout=1):
            is_signup = True
            if enable_register_log: print(f"[注册][{self.thread_id}] 检测到需要注册新账号")
            # 点击注册按钮
            try:
                sign_up_button = tab.ele("xpath=//button[contains(text(), '注册')]", timeout=1)
                if not sign_up_button:
                    sign_up_button = tab.ele("xpath=//button[contains(text(), 'Sign up')]", timeout=1)
                
                if sign_up_button:
                    sign_up_button.click()
                    if tab.wait.url_change(self.CURSOR_SIGNUP_URL, timeout=3):
                        print(f"[注册][{self.thread_id}] 进入注册页面")
            except Exception as e:
                print(f"[注册][{self.thread_id}] 点击注册按钮出错: {str(e)}")

        # 在密码页面选择验证方式或设置密码
        for retry in range(self.retry_times):
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 处理密码页面")
                
                # 检查是否在注册页面
                if is_signup or "sign-up" in tab.url:
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 设置密码: {password}")
                    # 设置新密码
                    password_input = tab.ele("xpath=//input[@name='password']", timeout=3)
                    if password_input:
                        password_input.input(password, clear=True)
                        tab.ele("@type=submit").click()
                else:
                    # 登录页面，使用邮件验证码
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 选择邮件验证码登录")
                    tab.ele("xpath=//button[@value='magic-code']").click()

                # 如果没有进入验证码页面，尝试通过验证码
                if not tab.wait.url_change(self.CURSOR_MAGAIC_CODE_URL, timeout=3) and \
                   not tab.wait.url_change(self.CURSOR_EMAIL_VERIFICATION_URL, timeout=3) and \
                   (self.CURSOR_PASSWORD_URL in tab.url or self.CURSOR_SIGNUP_PASSWORD_URL in tab.url):
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 尝试通过密码页面的验证码")
                    self._cursor_turnstile(tab)

            except Exception as e:
                print(f"[注册][{self.thread_id}] 处理密码页面时出现异常")
                print(e)

            # 检查是否进入验证码页面
            verification_url_changed = tab.wait.url_change(self.CURSOR_MAGAIC_CODE_URL, timeout=5) or \
                                       tab.wait.url_change(self.CURSOR_EMAIL_VERIFICATION_URL, timeout=5)
            if verification_url_changed:
                print(f"[注册][{self.thread_id}] 继续进入验证码页面")
                break

            # 检查是否有错误提示
            if tab.wait.eles_loaded("xpath=//p[contains(text(), 'Authentication blocked')]", timeout=3):
                print(f"[注册][{self.thread_id}][错误] 认证被阻止，请联系管理员")
                return tab, False

            if tab.wait.eles_loaded("xpath=//div[contains(text(), 'Sign up is restricted')]", timeout=3):
                print(f"[注册][{self.thread_id}][错误] 注册功能已被限制")
                return tab, False

            tab.refresh()
            # 超时退出函数
            if retry == self.retry_times - 1:
                if enable_register_log: print(f"[注册][{self.thread_id}] 处理密码页面超时")
                return tab, False

        # 获取邮件验证码
        try:
            verify_code = None

            print(f"[注册][{self.thread_id}] 等待接收验证码邮件...")
            data = self.email_queue.get(timeout=60)
            assert data is not None, "无法从邮件获取验证码"

            verify_code = self.parse_cursor_verification_code(data)
            assert verify_code is not None, "无法从邮件内容解析验证码"
            print(f"[注册][{self.thread_id}] 成功获取验证码: {verify_code}")
        except Exception as e:
            print(f"[注册][{self.thread_id}] 获取邮件验证码失败: {str(e)}")
            return tab, False

        # 输入邮件验证码
        for retry in range(self.retry_times):
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 输入邮件验证码")

                for idx, digit in enumerate(verify_code, start = 0):
                    tab.ele(f"xpath=//input[@data-index={idx}]").input(digit, clear=True)
                    tab.wait(0.1, 0.3)
                tab.wait(0.5, 1.5)

                # 如果没有进入主页，尝试通过验证码
                if not tab.wait.url_change(self.CURSOR_URL, timeout=3) and \
                   (self.CURSOR_MAGAIC_CODE_URL in tab.url or self.CURSOR_EMAIL_VERIFICATION_URL in tab.url):
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 尝试通过验证码页面的验证码")
                    self._cursor_turnstile(tab)

            except Exception as e:
                print(f"[注册][{self.thread_id}] 处理验证码页面时出现异常")
                print(e)

            # 检查是否进入主页
            if tab.wait.url_change(self.CURSOR_URL, timeout=3):
                print(f"[注册][{self.thread_id}] 账号注册/登录成功！")
                break

            tab.refresh()
            # 超时退出函数
            if retry == self.retry_times - 1:
                if enable_register_log: print(f"[注册][{self.thread_id}] 输入邮件验证码超时")
                return tab, False

        return tab, True
        
    def sign_up(self, email, password = None):
        """注册新账号"""
        assert self.email_server is not None, "应提供电子邮件服务器。"
 
        if self.email_server is not None:
            self.email_thread = threading.Thread(target=self.email_server.wait_for_new_message_thread,
                                                 args=(self.email_queue, ), 
                                                 daemon=True)
            self.email_thread.start()

        if password is None:
            fake = Faker()
            password = fake.password(length=12, special_chars=True, digits=True, upper_case=True, lower_case=True)

        tab = self.browser.new_tab(self.CURSOR_SIGNUP_URL)
        # 输入邮箱
        for retry in range(self.retry_times):
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 输入邮箱")
                tab.ele("xpath=//input[@name='email']").input(email, clear=True)
                tab.ele("@type=submit").click()

                # 如果没有进入密码页面，尝试通过验证码
                if not tab.wait.url_change(self.CURSOR_SIGNUP_PASSWORD_URL, timeout=3) and self.CURSOR_SIGNUP_URL in tab.url:
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 尝试通过邮箱页面的验证码")
                    self._cursor_turnstile(tab)

            except Exception as e:
                print(f"[注册][{self.thread_id}] 处理邮箱页面时出现异常")
                print(e)

            # 在密码页面或数据已验证，继续下一页
            if tab.wait.url_change(self.CURSOR_SIGNUP_PASSWORD_URL, timeout=5):
                print(f"[注册][{self.thread_id}] 继续进入密码页面")
                break

            tab.refresh()
            # 超时退出函数
            if retry == self.retry_times - 1:
                print(f"[注册][{self.thread_id}] 输入邮箱地址超时")
                return tab, False

        # 输入密码
        for retry in range(self.retry_times):
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 输入密码")
                tab.ele("xpath=//input[@name='password']").input(password, clear=True)
                tab.ele('@type=submit').click()

                # 如果没有进入验证码页面，尝试通过验证码
                if not tab.wait.url_change(self.CURSOR_EMAIL_VERIFICATION_URL, timeout=3) and self.CURSOR_SIGNUP_PASSWORD_URL in tab.url:
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 尝试通过密码页面的验证码")
                    self._cursor_turnstile(tab)

            except Exception as e:
                print(f"[注册][{self.thread_id}] 处理密码页面时出现异常")
                print(e)

            # 在验证码页面或数据已验证，继续下一页
            if tab.wait.url_change(self.CURSOR_EMAIL_VERIFICATION_URL, timeout=5):
                print(f"[注册][{self.thread_id}] 继续进入验证码页面")
                break

            if tab.wait.eles_loaded("xpath=//div[contains(text(), 'Sign up is restricted.')]", timeout=3):
                print(f"[注册][{self.thread_id}][错误] 注册功能已被限制")
                return tab, False

            tab.refresh()
            # 超时退出函数
            if retry == self.retry_times - 1:
                if enable_register_log: print(f"[注册][{self.thread_id}] 输入密码超时")
                return tab, False

        # 获取邮件验证码
        try:
            data = self.email_queue.get(timeout=60)
            assert data is not None, "无法从邮件获取验证码"

            verify_code = None
            if "body_text" in data:
                message_text = data["body_text"]
                message_text = message_text.replace(" ", "")
                verify_code = re.search(r'(?:\r?\n)(\d{6})(?:\r?\n)', message_text).group(1)
            elif "preview" in data:
                message_text = data["preview"]
                verify_code = re.search(r'Your verification code is (\d{6})\. This code expires', message_text).group(1)
            # 处理HTML格式
            elif "content" in data:
                message_text = data["content"]
                message_text = re.sub(r"<[^>]*>", "", message_text)
                message_text = re.sub(r"&#8202;", "", message_text)
                message_text = re.sub(r"&nbsp;", "", message_text)
                message_text = re.sub(r'[\n\r\s]', "", message_text)
                verify_code = re.search(r'openbrowserwindow\.(\d{6})Thiscodeexpires', message_text).group(1)
            assert verify_code is not None, "无法从邮件内容解析验证码"

        except Exception as e:
            print(f"[注册][{self.thread_id}] 获取邮件验证码失败: {str(e)}")
            return tab, False

        # 输入邮件验证码
        for retry in range(self.retry_times):
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 输入邮件验证码")

                for idx, digit in enumerate(verify_code, start = 0):
                    tab.ele(f"xpath=//input[@data-index={idx}]").input(digit, clear=True)
                    tab.wait(0.1, 0.3)
                tab.wait(0.5, 1.5)

                # 如果没有进入主页，尝试通过验证码
                if not tab.wait.url_change(self.CURSOR_URL, timeout=3) and self.CURSOR_EMAIL_VERIFICATION_URL in tab.url:
                    if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 尝试通过验证码页面的验证码")
                    self._cursor_turnstile(tab)

            except Exception as e:
                print(f"[注册][{self.thread_id}] 处理验证码页面时出现异常")
                print(e)

            # 检查是否进入主页
            if tab.wait.url_change(self.CURSOR_URL, timeout=3):
                print(f"[注册][{self.thread_id}] 账号注册成功！")
                break

            tab.refresh()
            # 超时退出函数
            if retry == self.retry_times - 1:
                if enable_register_log: print(f"[注册][{self.thread_id}] 输入邮件验证码超时")
                return tab, False

        return tab, True
    
    def get_usage(self, user_id):
        """获取账号使用情况"""
        tab = self.browser.new_tab(f"{self.CURSOR_USAGE_URL}?user={user_id}")
        return tab.json

    def delete_account(self):
        """删除账号"""
        tab = self.browser.new_tab(self.CURSOR_SETTING_URL)
        tab.ele("xpath=//div[contains(text(), 'Advanced')]").click()
        tab.ele("xpath=//button[contains(text(), 'Delete Account')]").click()
        tab.ele("""xpath=//input[@placeholder="Type 'Delete' to confirm"]""").input("Delete", clear=True)
        tab.ele("xpath=//span[contains(text(), 'Delete')]").click()
        return tab

    def parse_cursor_verification_code(self, email_data):
        """从邮件内容中解析验证码"""
        message = ""
        verify_code = None

        if "content" in email_data:
            message = email_data["content"]
            
            # 尝试提取 HTML 格式的验证码 (<div style="font-family:-apple-system,BlinkMacSystemFont,...">123456</div>)
            html_code_match = re.search(r'<div style=".*?">(\d{6})</div>', message)
            if html_code_match:
                verify_code = html_code_match.group(1)
                print(f"[注册] 找到HTML格式验证码: {verify_code}")
                return verify_code

            # 尝试提取 aria-label 中的验证码
            aria_label_match = re.search(r'aria-label="[^"]*Your one-time code is (\d{6})[^"]*"', message)
            if aria_label_match:
                verify_code = aria_label_match.group(1)
                print(f"[注册] 找到aria-label验证码: {verify_code}")
                return verify_code
            
            # 尝试任何出现的 6 位数字
            any_code_match = re.search(r'\b(\d{6})\b', message)
            if any_code_match:
                verify_code = any_code_match.group(1)
                print(f"[注册] 找到通用验证码: {verify_code}")
                return verify_code
            
            # 传统的提取方法
            message = message.replace(" ", "")
            try:
                verify_code = re.search(r'(?:\r?\n)(\d{6})(?:\r?\n)', message).group(1)
                print(f"[注册] 找到传统格式验证码: {verify_code}")
                return verify_code
            except:
                pass
                
        elif "text" in email_data:
            message = email_data["text"]
            message = message.replace(" ", "")
            try:
                verify_code = re.search(r'(?:\r?\n)(\d{6})(?:\r?\n)', message).group(1)
                print(f"[注册] 找到文本格式验证码: {verify_code}")
                return verify_code
            except:
                # 尝试任何出现的 6 位数字
                any_code_match = re.search(r'\b(\d{6})\b', message)
                if any_code_match:
                    verify_code = any_code_match.group(1)
                    print(f"[注册] 找到通用文本验证码: {verify_code}")
                    return verify_code
        
        print(f"[注册] 无法在邮件中找到验证码")
        return verify_code

    def get_cursor_cookie(self, tab):
        """获取Cursor账号的Cookie和Token"""
        def _generate_pkce_pair():
            code_verifier = secrets.token_urlsafe(43)
            code_challenge_digest = hashlib.sha256(code_verifier.encode('utf-8')).digest()
            code_challenge = base64.urlsafe_b64encode(code_challenge_digest).decode('utf-8').rstrip('=')    
            return code_verifier, code_challenge
        try:
            verifier, challenge = _generate_pkce_pair()
            id = uuid.uuid4()
            client_login_url = f"https://www.cursor.com/cn/loginDeepControl?challenge={challenge}&uuid={id}&mode=login"
            tab.get(client_login_url)
            tab.ele("xpath=//span[contains(text(), 'Yes, Log In')]").click()

            auth_pooll_url = f"https://api2.cursor.sh/auth/poll?uuid={id}&verifier={verifier}"
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Cursor/0.48.6 Chrome/132.0.6834.210 Electron/34.3.4 Safari/537.36",
                "Accept": "*/*"
            }
            response = requests.get(auth_pooll_url, headers = headers, timeout=5)
            data = response.json()
            accessToken = data.get("accessToken", None)
            authId = data.get("authId", "")
            if len(authId.split("|")) > 1:
                userId = authId.split("|")[1]
                token = f"{userId}%3A%3A{accessToken}"
            else:
                token = accessToken
        except Exception as e:
            print(f"[注册][{self.thread_id}] 获取Cookie失败: {str(e)}")
            return None

        if enable_register_log:
            if token is not None:
                print(f"[注册][{self.thread_id}] 成功获取账号Cookie和Token")
            else:
                print(f"[注册][{self.thread_id}] 获取账号Cookie和Token失败")
        return token

    def _cursor_turnstile(self, tab, retry_times = 5):
        """处理Turnstile验证码"""
        for retry in range(retry_times): 
            try:
                if enable_register_log: print(f"[注册][{self.thread_id}][{retry}] 正在通过Turnstile验证")
                challenge_shadow_root = tab.ele('@id=cf-turnstile').child().shadow_root
                challenge_shadow_button = challenge_shadow_root.ele("tag:iframe", timeout=30).ele("tag:body").sr("xpath=//input[@type='checkbox']")
                if challenge_shadow_button:
                    challenge_shadow_button.click()
                    break
            except:
                pass
            if retry == retry_times - 1:
                print("[注册] 通过验证码超时")
